{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/mlop-ai/mlop/blob/main/examples/lightning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "<h1 align=\"center\" style=\"font-family: Inter, sans-serif; font-style: normal; font-weight: 700; font-size: 72px\">m:lop</h1>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dU4sYW0PNMNe",
    "outputId": "ba2a6082-8c7b-4997-864c-bd281ee0f8dc"
   },
   "outputs": [],
   "source": [
    "%pip install -Uq \"mlop[full]\" lightning torchmetrics onnx\n",
    "# %pip install \"mlop[full] @ git+https://github.com/mlop-ai/mlop.git\"\n",
    "# import sys; import os; sys.path.insert(0, os.path.dirname(os.path.abspath(os.path.dirname(\"__file__\"))))\n",
    "import mlop\n",
    "from mlop.compat.lightning import MLOPLogger\n",
    "\n",
    "mlop.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7qJSFFVoNMNf"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision import transforms\n",
    "import lightning.pytorch as pl\n",
    "import torchmetrics\n",
    "\n",
    "pl.seed_everything(1337)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the experiment run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BIjra-peNMNj"
   },
   "outputs": [],
   "source": [
    "class LitMLP(pl.LightningModule):\n",
    "    def __init__(self, in_dims, n_classes=10, n_layer_1=128, n_layer_2=256, lr=1e-4):\n",
    "        super().__init__()\n",
    "\n",
    "        self.layer_1 = nn.Linear(np.prod(in_dims), n_layer_1)\n",
    "        self.layer_2 = nn.Linear(n_layer_1, n_layer_2)\n",
    "        self.layer_3 = nn.Linear(n_layer_2, n_classes)\n",
    "\n",
    "        self.save_hyperparameters()\n",
    "        self.train_acc = torchmetrics.Accuracy(task=\"multiclass\", num_classes=n_classes)\n",
    "        self.valid_acc = torchmetrics.Accuracy(task=\"multiclass\", num_classes=n_classes)\n",
    "        self.test_acc = torchmetrics.Accuracy(task=\"multiclass\", num_classes=n_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size, *dims = x.size()\n",
    "\n",
    "        x = x.view(batch_size, -1)\n",
    "        x = F.relu(self.layer_1(x))\n",
    "        x = F.relu(self.layer_2(x))\n",
    "        x = self.layer_3(x)\n",
    "        x = F.log_softmax(x, dim=1)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def loss(self, xs, ys):\n",
    "        logits = self(xs)\n",
    "        loss = F.nll_loss(logits, ys)\n",
    "        return logits, loss\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        xs, ys = batch\n",
    "        logits, loss = self.loss(xs, ys)\n",
    "        preds = torch.argmax(logits, 1)\n",
    "\n",
    "        self.log(\"train/loss\", loss, on_epoch=True)\n",
    "        self.train_acc(preds, ys)\n",
    "        self.log(\"train/acc\", self.train_acc, on_epoch=True)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(self.parameters(), lr=self.hparams[\"lr\"])\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        xs, ys = batch\n",
    "        logits, loss = self.loss(xs, ys)\n",
    "        preds = torch.argmax(logits, 1)\n",
    "\n",
    "        self.test_acc(preds, ys)\n",
    "        self.log(\"test/loss_epoch\", loss, on_step=False, on_epoch=True)\n",
    "        self.log(\"test/acc_epoch\", self.test_acc, on_step=False, on_epoch=True)\n",
    "\n",
    "    def on_test_epoch_end(self):  # args are defined as part of pl API\n",
    "        dummy_input = torch.zeros(self.hparams[\"in_dims\"], device=self.device)\n",
    "        model_filename = \"model_final.onnx\"\n",
    "        self.to_onnx(model_filename, dummy_input, export_params=True)\n",
    "\n",
    "        # TODO: remove legacy compat\n",
    "        artifact = mlop.Artifact(name=\"model.ckpt\", type=\"model\")\n",
    "        artifact.add_file(model_filename)\n",
    "        self.logger.experiment.log({\"model\": artifact})\n",
    "\n",
    "    def on_validation_epoch_start(self):\n",
    "        self.validation_step_outputs = []\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        xs, ys = batch\n",
    "        logits, loss = self.loss(xs, ys)\n",
    "        preds = torch.argmax(logits, 1)\n",
    "        self.valid_acc(preds, ys)\n",
    "\n",
    "        self.log(\"valid/loss_epoch\", loss)  # default on val/test is on_epoch only\n",
    "        self.log(\"valid/acc_epoch\", self.valid_acc)\n",
    "        self.validation_step_outputs.append(logits)\n",
    "\n",
    "        return logits\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        validation_step_outputs = self.validation_step_outputs\n",
    "        dummy_input = torch.zeros(self.hparams[\"in_dims\"], device=self.device)\n",
    "        model_filename = f\"model_{str(self.global_step).zfill(5)}.onnx\"\n",
    "        torch.onnx.export(self, dummy_input, model_filename, opset_version=11)\n",
    "\n",
    "        # TODO: remove legacy compat\n",
    "        artifact = mlop.Artifact(name=\"model.ckpt\", type=\"model\")\n",
    "        artifact.add_file(model_filename)\n",
    "        self.logger.experiment.log({\"model\": artifact})\n",
    "        \n",
    "        flattened_logits = torch.flatten(torch.cat(validation_step_outputs))\n",
    "        self.logger.experiment.log(\n",
    "            {\n",
    "                \"valid/logits\": mlop.Histogram(flattened_logits.to(\"cpu\")),\n",
    "                \"global_step\": self.global_step,\n",
    "            }\n",
    "        )\n",
    "\n",
    "\n",
    "class ImagePredictionLogger(pl.Callback):\n",
    "    def __init__(self, val_samples, num_samples=32):\n",
    "        super().__init__()\n",
    "        self.val_imgs, self.val_labels = val_samples\n",
    "        self.val_imgs = self.val_imgs[:num_samples]\n",
    "        self.val_labels = self.val_labels[:num_samples]\n",
    "\n",
    "    def on_validation_epoch_end(self, trainer, pl_module):\n",
    "        val_imgs = self.val_imgs.to(device=pl_module.device)\n",
    "        logits = pl_module(val_imgs)\n",
    "        preds = torch.argmax(logits, 1)\n",
    "        trainer.logger.experiment.log(\n",
    "            {\n",
    "                \"examples\": [\n",
    "                    mlop.Image(x, caption=f\"Pred:{pred}, Label:{y}\")\n",
    "                    for x, pred, y in zip(val_imgs, preds, self.val_labels)\n",
    "                ],\n",
    "                \"global_step\": trainer.global_step,\n",
    "            }\n",
    "        )\n",
    "\n",
    "\n",
    "class MNISTDataModule(pl.LightningDataModule):\n",
    "    def __init__(self, data_dir=\".lightning/\", batch_size=128):\n",
    "        super().__init__()\n",
    "        self.data_dir = data_dir\n",
    "        self.batch_size = batch_size\n",
    "        self.transform = transforms.Compose(\n",
    "            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
    "        )\n",
    "\n",
    "    def prepare_data(self):\n",
    "        MNIST(self.data_dir, train=True, download=True)\n",
    "        MNIST(self.data_dir, train=False, download=True)\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        if stage == \"fit\" or stage is None:\n",
    "            mnist = MNIST(self.data_dir, train=True, transform=self.transform)\n",
    "            self.mnist_train, self.mnist_val = random_split(mnist, [55000, 5000])\n",
    "        if stage == \"test\" or stage is None:\n",
    "            self.mnist_test = MNIST(\n",
    "                self.data_dir, train=False, transform=self.transform\n",
    "            )\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        mnist_train = DataLoader(self.mnist_train, batch_size=self.batch_size)\n",
    "        return mnist_train\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        mnist_val = DataLoader(self.mnist_val, batch_size=10 * self.batch_size)\n",
    "        return mnist_val\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        mnist_test = DataLoader(self.mnist_test, batch_size=10 * self.batch_size)\n",
    "        return mnist_test\n",
    "\n",
    "\n",
    "mnist = MNISTDataModule()\n",
    "mnist.prepare_data()\n",
    "mnist.setup()\n",
    "samples = next(iter(mnist.val_dataloader()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start training with **mlop**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TaijaXgRNMNs"
   },
   "outputs": [],
   "source": [
    "mlop_logger = MLOPLogger(name=\".lightning\")\n",
    "trainer = pl.Trainer(\n",
    "    logger=mlop_logger,\n",
    "    log_every_n_steps=50,\n",
    "    max_epochs=5,\n",
    "    deterministic=True,\n",
    "    callbacks=[ImagePredictionLogger(samples)],\n",
    ")\n",
    "\n",
    "model = LitMLP(in_dims=(1, 28, 28))\n",
    "try:\n",
    "    trainer.fit(model, mnist)\n",
    "    trainer.test(datamodule=mnist, ckpt_path=None)\n",
    "finally:\n",
    "    mlop.finish()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
